import os

import numpy
import pandas as pd
import random
import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

tp = transforms.ToPILImage()
tt = transforms.ToTensor()


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    numpy.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True


def norm_imgt(imgt: torch.Tensor):
    return (imgt - imgt.min()) / (imgt.max() - imgt.min())


def preprocess(label_path, img_path, out_path, offset=0):
    df = pd.read_csv(label_path)
    imgs = []
    for i in range(df.shape[0]):
        im = Image.open(img_path + '{:0>3d}.bmp'.format(df['image ID'][i])).convert('RGB')
        seg_im = Image.open(img_path + '{:0>3d}.png'.format(df['image ID'][i]))
        im = im.resize((128, 128))
        seg_im = seg_im.resize((128, 128))
        im.save(out_path + '{:0>3d}.bmp'.format(df['image ID'][i] + offset))
        seg_im.save(out_path + '{:0>3d}.png'.format(df['image ID'][i] + offset))


def merge_datasets():
    os.mkdir('./segmentation_WBC-master/merged_dataset/')

    df1 = pd.read_csv('./segmentation_WBC-master/Class Labels of Dataset 1.csv')
    df2 = pd.read_csv('./segmentation_WBC-master/Class Labels of Dataset 2.csv')
    df2['image ID'] += 300

    merged_df = pd.concat([df1, df2])
    merged_df.to_csv('./segmentation_WBC-master/merged_dataset/merged_labels.csv', index=False)

    os.system('mv ./segmentation_WBC-master/preproc_ds1/* ./segmentation_WBC-master/merged_dataset/')
    os.system('mv ./segmentation_WBC-master/preproc_ds2/* ./segmentation_WBC-master/merged_dataset/')


# %%
class MyDataset(Dataset):
    def __init__(self, img_path, label_path, transform=None):
        df = pd.read_csv(label_path)
        imgs = []
        for i in range(df.shape[0]):
            imgs.append((img_path + '{:0>3d}.bmp'.format(df['image ID'][i]),
                         img_path + '{:0>3d}.png'.format(df['image ID'][i]),
                         df['class'][i] - 1))
        self.imgs = imgs
        self.transform = transform

    def __getitem__(self, index):
        fn1, fn2, label = self.imgs[index]
        img = Image.open(fn1).convert('RGB')
        seg_img = Image.open(fn2)
        if self.transform is not None:
            img = self.transform(img)
        if self.transform is not None:
            seg_img = self.transform(seg_img)
        return img, seg_img, label

    def __len__(self):
        return len(self.imgs)


class MyDataset2(Dataset):
    def __init__(self, img_path, label_path, transform=None, target_transform=None):
        df = pd.read_csv(label_path)
        imgs = []
        for i in range(df.shape[0]):
            imgs.append((img_path + '{:0>3d}.png'.format(df['image ID'][i]),
                         int(df['class'][i])))
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):
        fn, label = self.imgs[index]
        seg_img = Image.open(fn)
        if self.transform is not None:
            seg_img = self.transform(seg_img)
        if self.target_transform is not None:
            label = self.target_transform(torch.tensor(label))
        return seg_img, label

    def __len__(self):
        return len(self.imgs)


def intermediate_result(net, dataset: Dataset, out_path, gt_path, new_labels, dev='cpu'):
    df = pd.DataFrame(columns=('image ID', 'class'))
    for (idx, dt) in enumerate(dataset):
        out_tensor = net(dt[0].reshape((1, *dt[0].shape)).to(dev))
        out_tensor = norm_imgt(out_tensor[0].cpu())
        out_img = tp(out_tensor)
        out_img.save(out_path + '{:0>3d}.png'.format(idx + 1))
        gt_img = tp(dt[1])
        gt_img.save(gt_path + '{:0>3d}.png'.format(idx + 1))
        s = pd.Series({'image ID': idx + 1, 'class': dt[2]})
        df = df.append(s, ignore_index=True)
    df.to_csv(new_labels, index=False)


def classification_results(net, dataset: Dataset, outpath, dev='cpu'):
    df = pd.DataFrame(columns=('yhat', 'y'))
    for (idx, dt) in enumerate(dataset):
        out_tensor = net(dt[0].reshape((1, *dt[0].shape)).to(dev))
        cls = torch.argmax(out_tensor).cpu().item() + 1
        s = pd.Series({'yhat': cls, 'y': dt[1] + 1})
        df = df.append(s, ignore_index=True)
    df.to_csv(outpath, index=False)
